#!/bin/bash

# LayerNorm拆分脚本
# 用于拆分ONNX模型中的LayerNorm算子为基础算子

echo "=== ONNX LayerNorm算子拆分工具 ==="

# 检查输入文件
INPUT_FILE="co_detr_optimized.onnx"
if [ ! -f "$INPUT_FILE" ]; then
    echo "错误: 输入文件不存在: $INPUT_FILE"
    echo "请确保 co_detr_optimized.onnx 文件在当前目录中"
    exit 1
fi

echo "输入文件: $INPUT_FILE"
echo "文件大小: $(du -h $INPUT_FILE | cut -f1)"

# 检查Python环境
if ! command -v python3 &> /dev/null; then
    echo "错误: 未找到 python3"
    exit 1
fi

# 检查必要的Python包
echo "检查Python依赖..."
python3 -c "import onnx" 2>/dev/null
if [ $? -ne 0 ]; then
    echo "错误: 未安装 onnx 包"
    echo "请运行: pip install onnx"
    exit 1
fi

python3 -c "import numpy" 2>/dev/null
if [ $? -ne 0 ]; then
    echo "错误: 未安装 numpy 包"
    echo "请运行: pip install numpy"
    exit 1
fi

echo "✓ Python环境检查通过"

# 选择拆分方式
echo ""
echo "请选择拆分方式:"
echo "1) 基础拆分 (split_layernorm.py)"
echo "2) 高级拆分 (split_layernorm_advanced.py)"
echo "3) 仅分析模型结构"
echo "4) 退出"
echo ""
read -p "请输入选择 (1-4): " choice

case $choice in
    1)
        echo "执行基础LayerNorm拆分..."
        OUTPUT_FILE="co_detr_layernorm_split.onnx"
        python3 split_layernorm.py --input "$INPUT_FILE" --output "$OUTPUT_FILE"
        ;;
    2)
        echo "执行高级LayerNorm拆分..."
        OUTPUT_FILE="co_detr_layernorm_split_advanced.onnx"
        python3 split_layernorm_advanced.py --input "$INPUT_FILE" --output "$OUTPUT_FILE"
        ;;
    3)
        echo "分析模型结构..."
        python3 split_layernorm_advanced.py --input "$INPUT_FILE" --analyze-only
        ;;
    4)
        echo "退出"
        exit 0
        ;;
    *)
        echo "无效选择"
        exit 1
        ;;
esac

# 检查输出文件
if [ -f "$OUTPUT_FILE" ]; then
    echo ""
    echo "=== 拆分结果 ==="
    echo "输出文件: $OUTPUT_FILE"
    echo "文件大小: $(du -h $OUTPUT_FILE | cut -f1)"
    echo "原始文件大小: $(du -h $INPUT_FILE | cut -f1)"
    
    # 比较文件大小
    ORIGINAL_SIZE=$(stat -c%s "$INPUT_FILE")
    OUTPUT_SIZE=$(stat -c%s "$OUTPUT_FILE")
    SIZE_RATIO=$(echo "scale=2; $OUTPUT_SIZE / $ORIGINAL_SIZE" | bc)
    
    echo "大小比例: ${SIZE_RATIO}x"
    
    if (( $(echo "$SIZE_RATIO > 1.1" | bc -l) )); then
        echo "⚠ 新模型比原始模型大，可能包含更多基础算子"
    elif (( $(echo "$SIZE_RATIO < 0.9" | bc -l) )); then
        echo "✓ 新模型比原始模型小"
    else
        echo "✓ 新模型大小与原始模型相近"
    fi
    
    echo ""
    echo "✓ LayerNorm拆分完成!"
    echo "新模型已保存为: $OUTPUT_FILE"
else
    echo "⚠ 未生成输出文件，请检查错误信息"
fi 